% Chem-Lab! 用于模拟简单的基元反应，支持多种物质和化学反应式（至少希望是！）
% 比如在基元反应 aA+bB->cC中，
% 正向化学反应(aA+bB->cC)速率  vp = kp [A]^a [B]^b （以化学反应进度记）
% 逆向化学反应(cC->aA+bB)速率  vr = kr [C]^c
% 那么 d[A]/dt = - a vp + a vr， d[C]/dt = + c vp - c vr 等
% 此外，当反应达到平衡时，vp=vr => = kp [A]^a [B]^b = kr [C]^c
% => kp/kr =[C]^c /([A]^a [B]^b) = K
% 即 kp/kr = K, K 是反应平衡常数
% 可能有bug，仅供娱乐，请勿用于正式用途！
% Gitee Repo

clc
clear

global N
global N_eq

%Time step
dt = 0.001;
%Maximum tick
TICK = 20000;

%Num of Compounds
N = 4;
%Num of Chemical Equations
N_eq = 2;

% Spatial Gird and Compounds
L = 5;
dx = 0.05;
[x,y] = meshgrid(-L:dx:L);
n = size(x,1);
u0 = zeros(n,n,N);
u1 = zeros(n,n,N);
D = ones(N,1);
D = [0.5;0.5;0.5;0.5];

u0(:,:,4) = exp(-5*((x).^2+y.^2));

% Chemical Equations
% Eq.nu, Eq.kp, Eq.kr.

% Eq.nu: 化学计量数
% -: Reactant, +: Production,0: Not participate in this reaction
% Exp. A+B->C, [-1,-1,1,0]
Eq(1).nu =[-1,-1,1,0];
%Eq.kp : 正向反应速率系数 A+B->C
Eq(1).kp = 1;
%Eq.kr : 逆向反应速率系数 C->A+B
Eq(1).kr = 10;
% 正向反应速率系数与逆向反应速率系数由反应平衡常数联系
% kp/kr = K, K is the equilibrium constant

% Exp. C->D, [0,0,-1,1]
Eq(2).nu =[0,0,-1,1];
Eq(2).kp = 100;
Eq(2).kr = 1;
% 这个例子展示了一个包括反应中间体的例子
% 该反应由两步组成, A+B->C, C->D, C是反应中间体
% C很难被生成，也很快会分解。尽管C的含量一直很少，但是C是反应的必要一环

% Chek whether equation is valid
for j = 1:N_eq
    if !any(Eq(j).nu >0) || !any(Eq(j).nu < 0) || all(Eq(j).nu == 0)
        error('WRONG EQ\n');
        return;
    end
end

%Compute Reaction Rate
% u (n,n,N)
% eq is a struct
% vp: Rate of positive Reaction
% vr: Rate of negative Reaction
function [vp, vr] = compute_v(u, eq)
    global N
    global N_eq

    vp = eq.kp;
    vr = eq.kr;

    for i = 1:N
        if eq.nu(i) >0
            vr = vr.*u(:,:,i).^(eq.nu(i));
        elseif  eq.nu(i) == 0
            % Not participate in this reaction
            ;
        else
            vp = vp.*u(:,:,i).^(abs(eq.nu(i)));
        end
    end
    return
end


con = zeros(TICK,N);

% Main iteration
figure()
for tick = 1:TICK
    if any(u0<0) || any(isnan(u0)) || any(isinf(u0))
        error('Substance consumed. Time step too long?\n');
    end

    u1 = u0 + dt/(dx).^2*reshape(D,1,1,N).*(circshift(u0,-1,1) -2*u0+circshift(u0,1,1) + circshift(u0,-1,2) -2*u0 + circshift(u0,1,2));
    u0 = u1;

    u1 = u0;
    for j = 1:N_eq
        [vp, vr] = compute_v(u0, Eq(j));
        u1 += reshape(Eq(j).nu,1,1,N).*(vp-vr)*dt;
    end
    u0 = u1;

    for i = 1:N
        con(tick,i) = sum(u1(:,:,i)(:));
    end

    % Print debug info
    if mod(tick,500) == 0
        clf
        hold on
        axis equal
        xlabel('x')
        ylabel('y')
        zlabel('u')

        labels = cell(1, N);
        for i = 1:N
            surf(x,y,u1(:,:,i),'Edgecolor','None')
            labels{i} = num2str(i);
            disp(sum(u1(:,:,i)(:)))
        end
        legend(labels)
        drawnow
        pause(0.1)
    end
end

figure
hold on
##axis equal
_t = (1:TICK)';
for i = 1:N
    plot(_t,con(:,i));
end

